import asyncio
import datetime
import json
import logging
import random
import os
import shutil
import time
import uuid
import argparse
from concurrent.futures import ProcessPoolExecutor

import cv2
import aiofiles
import aiohttp
import websockets

script_help = "有时压测脚本, 使用示例：python youshi_test.py --workers 100 --http-server " \
              "https://growup-test.aminer.cn/api --ws-server wss://growup-test.aminer.cn/api " \
              "--interval 1"
parser = argparse.ArgumentParser(description=script_help)
http_server_help = "http接口的请求地址，需包含整个地址的前缀，如：https://growup.aminer.cn/prod/api/peilian"
parser.add_argument('--http-server', type=str, default="https://growup-test.aminer.cn/api", help=http_server_help)
ws_server_help = "websocket连接的地址信息, 需包含整个地址的前缀，如：ws://growup.aminer.cn/prod/api/peilian"
parser.add_argument('--ws-server', type=str, default="wss://growup-test.aminer.cn/api", help=ws_server_help)
parser.add_argument("--workers", type=int, default=1, help='并发数量')
interval_help = '对于给定的并发量，多长时间启动一个新的任务, 如：若该值为1,表示每隔1秒启动一场面试，启动100场面试需要100秒'
parser.add_argument("--interval", type=float, default=1, help=interval_help)
parser.add_argument("--process", type=int, default=1, help='进程数量,可利用多核CPU特性')
args = parser.parse_args()

LOG_LEVEL = logging.INFO
LOG_PATH = os.path.join(os.path.dirname(__file__), "logs")
VIDEO_PATH = os.path.join(os.path.dirname(__file__), "video")
PCM_DIR = os.path.join(os.path.dirname(__file__), 'pcm')

os.system(f'mkdir -p {LOG_PATH}')
os.system(f'mkdir -p {VIDEO_PATH}')

total_count = 0
lock = asyncio.Lock()


def clear_dir(path):
    for files in os.listdir(path):
        file_path = os.path.join(path, files)
        try:
            shutil.rmtree(file_path)
        except OSError:
            os.remove(file_path)


if __name__ == '__main__':
    clear_dir(LOG_PATH)
    clear_dir(VIDEO_PATH)


def get_logger(name):
    logger = logging.getLogger(name)
    filename = os.path.join(LOG_PATH, f'{name}.log')
    fh = logging.FileHandler(filename, mode='w+', encoding='utf-8')
    ch = logging.StreamHandler()
    formatter = logging.Formatter('%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s')
    logger.setLevel(LOG_LEVEL)
    fh.setFormatter(formatter)
    ch.setFormatter(formatter)
    logger.addHandler(fh)
    logger.addHandler(ch)
    return logger


class InterruptError(Exception):
    pass


# 用户登录
async def login(session) -> str:
    url = args.http_server + "/account/login"
    data = {
        "type": "phone",
        "ph": {
            "phone": "18435154915",
            "code": "123456"
        },
        "promote_id": ""
    }
    async with session.post(url, json=data) as resp:
        resp_json = await resp.json()
        print(resp_json)
        if resp_json.get('code') == 200:
            return resp_json['data']['token']
        raise InterruptError("登录失败")


# 获取岗位列表
async def question_groups(session) -> list[int]:
    url = args.http_server + "/question/groups"
    async with session.get(url) as resp:
        resp_json = await resp.json()
        if resp_json.get('code') == 200:
            career_id_li = list()
            for career_type in resp_json['data']:
                for career in career_type['value']:
                    career_id_li.append(career['id'])
            return career_id_li
        raise InterruptError("获取岗位列表失败")


# 创建面试场次，返回场次ID
async def create_practice(session, career_id, token) -> int:
    url = args.http_server + "/practice/create"
    data = {'career_id': career_id}
    headers = {'authenticationtoken': token}
    try:
        async with session.post(url, json=data, headers=headers) as resp:
            resp_json = await resp.json()
            if resp_json.get('code') == 200:
                return resp_json['data']['id']
            raise InterruptError("创建面试场次失败")
    except Exception as e:
        print(e)


async def ws_keepalive(ws, logger):
    while True:
        try:
            await ws.send("ping")
        except websockets.ConnectionClosed as e:
            logger.warning(f'[{logger.name}]{e}')
            break
        await asyncio.sleep(15)


def video_duration(video_path) -> float:
    video = cv2.VideoCapture(video_path)
    if video.isOpened():
        rate = video.get(5)
        frame_num = video.get(7)
        duration = frame_num / rate
        return duration
    return 3


async def download_video(url, video_name, logger):
    t1 = time.time()
    async with aiohttp.ClientSession() as s:
        async with s.get(url) as resp:
            if resp.status == 200:
                video_path = os.path.join(VIDEO_PATH, video_name)
                base_dir = os.path.dirname(video_path)
                os.system(f"mkdir -p {base_dir}")
                f = await aiofiles.open(video_path, mode="wb")
                await f.write(await resp.read())
                await f.close()
            else:
                logger.error(f'[{logger.name}]下载【{url}】失败')
                raise InterruptError(f'[{logger.name}]下载【{url}】失败')

    t2 = time.time()
    logger.info(f"[{logger.name}]视频下载使用了[{t2 - t1}s]")
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, video_duration, video_path)


async def send_pcm(ws, logger):
    pcm_files = [
        os.path.join(PCM_DIR, f) for f in os.listdir(PCM_DIR)
        if os.path.isfile(os.path.join(PCM_DIR, f))
    ]
    pcm_file = random.choice(pcm_files)
    logger.info(f"pcm_file:[{pcm_file}]")
    file_obj = await aiofiles.open(pcm_file, 'rb')
    try:
        index = 1
        while True:
            chunk = await file_obj.read(1280)
            if not chunk:
                break
            await ws.send(chunk)

            index += 1
            await asyncio.sleep(0.04)
    finally:
        await file_obj.close()


async def play_video(res_json, logger, session_id):
    video_url = res_json['data']['video']
    current_tag = res_json['data']['current']
    video_type = res_json['type']
    video_save_path = f"{session_id}/{video_type}_{current_tag}_{str(uuid.uuid4())[:5]}.mp4"
    duration = await download_video(video_url, video_save_path, logger)
    logger.info(f"[{session_id}]面试视频播放中[{duration}s]")
    duration = 10
    await asyncio.sleep(duration)


async def process_ws(ws, logger, session_id):
    await ws.send("AUDIO::INIT")
    logger.info(f"[{session_id}]开始答题")
    try:
        while True:
            res = await ws.recv()
            res_json = json.loads(res)
            if res_json['type'] != "pong":
                logger.info(f'[{session_id}]{res}')
            match res_json['type']:
                case "pong":
                    pass
                case "ask" | "extra":
                    await play_video(res_json, logger, session_id)
                    await ws.send("AUDIO::BEGIN")
                    logger.info(f"[{session_id}]开始回答（发送pcm文件）")
                    await send_pcm(ws, logger)
                    await ws.send("AUDIO::FINISH::35")
                    logger.info(f"[{session_id}]结束回答")
                case "conv":
                    await play_video(res_json, logger, session_id)
                    await ws.send("AUDIO::CONV::FINISH")
                case "finish":
                    await play_video(res_json, logger, session_id)
                    return
                case _:
                    logger.warning(f"[{session_id}]未知消息: {res}")
    except Exception as e:
        logger.error(f'[{session_id}]{e}')
        return e


async def websocket_conn(session_id, token, logger):
    ws_url = args.ws_server + f"/practice/session/{session_id}/ask"
    headers = {"authenticationtoken": token}
    async with websockets.connect(ws_url, extra_headers=headers, ping_timeout=60) as ws:
        asyncio.create_task(ws_keepalive(ws, logger))
        return await process_ws(ws, logger, session_id)


async def gen_session_report(session, session_id, token, logger):
    url = args.http_server + f'/practice/session/{session_id}/report'
    data = {'session_id': str(session_id)}
    headers = {'authenticationtoken': token}
    async with session.post(url, json=data, headers=headers) as resp:
        resp_json = await resp.json()
        logger.info(f'[{session_id}]{resp_json}')
        if resp_json.get("code") == 200:
            logger.info(f'[{session_id}]的报告正在生成, 请稍作等待')
        else:
            logger.warning(f'[{session_id}]的报告生成错误')


async def start():
    pressure_test_start = datetime.datetime.now()
    pressure_test_end = datetime.datetime.now()

    while pressure_test_end - pressure_test_start < datetime.timedelta(hours=2):
        async with aiohttp.ClientSession() as session:
            token = await login(session)
            career_id_li = await question_groups(session)
            career_id = random.choice(career_id_li)
            session_id = await create_practice(session, career_id, token)
            logger = get_logger(str(session_id))
            logger.info(args.http_server)
            logger.info(args.ws_server)
            logger.info(f'workers: [{args.workers}]')
            logger.info(f"token: {token}")
            logger.info(f"career_id: {career_id}")
            logger.info(f"session_id: {session_id}")
            result = await websocket_conn(session_id, token, logger)
            if result is not None:
                return
            await gen_session_report(session, session_id, token, logger)
            logger.info("面试完成")

        async with lock:
            global total_count
            total_count += 1
        await asyncio.sleep(2)
        pressure_test_end = datetime.datetime.now()


async def main():
    task_li = list()
    for _ in range(args.workers):
        task = asyncio.create_task(start())
        task_li.append(task)
        await asyncio.sleep(args.interval)

    await asyncio.wait(task_li)


def run():
    asyncio.run(main())


if __name__ == '__main__':
    with ProcessPoolExecutor(max_workers=args.process) as executor:
        for _ in range(args.process):
            executor.submit(run)
